#!/usr/bin/env python3
import rclpy
from rclpy.node import Node
from geometry_msgs.msg import TransformStamped
from tf2_ros import StaticTransformBroadcaster
import tf2_ros
from sensor_msgs.msg import Imu
import tf_transformations

class StaticTfPublisher(Node):
    def __init__(self):
        super().__init__('static_tf_publisher')
        # 创建 TF2 广播器
        self.tf_broadcaster = tf2_ros.TransformBroadcaster(self)

        # 订阅 IMU 数据
        self.subscription = self.create_subscription(
            Imu,
            '/imu',  # IMU 话题名称
            self.imu_callback,
            10
        )

        # 初始化变换
        self.transform_stamped = TransformStamped()
        self.transform_stamped.header.frame_id = "odom"  # 父坐标系
        self.transform_stamped.child_frame_id = "base_footprint"        # 子坐标系

        self.map_to_odom = TransformStamped()
        self.map_to_odom.header.frame_id = "map"
        self.map_to_odom.child_frame_id = "odom"

    def imu_callback(self, msg):
        # 从 IMU 数据中提取四元数
        orientation = msg.orientation
        q = [orientation.w, orientation.x, orientation.y, orientation.z]
 
        # 将四元数转换为欧拉角 (roll, pitch, yaw)
        (roll, pitch, yaw) = tf_transformations.euler_from_quaternion(q)
 
        # 创建一个新的四元数，仅包含围绕 Z 轴的旋转
        # 注意：这里我们实际上只关心 yaw，但变换消息需要完整的四元数
        # 为了简化，我们假设只围绕 Z 轴旋转，其他轴为 0
        # 这种情况下，四元数可以简化为围绕 Z 轴旋转的四元数
        # 使用 yaw 创建新的四元数
        q_z_axis = tf_transformations.quaternion_from_euler(roll, 0, 0)
 
        # 更新变换中的旋转
        self.transform_stamped.transform.rotation.x = q_z_axis[1]
        self.transform_stamped.transform.rotation.y = q_z_axis[2]
        self.transform_stamped.transform.rotation.z = q_z_axis[3]
        self.transform_stamped.transform.rotation.w = q_z_axis[0]
 
        # 设置时间戳和坐标系
        self.transform_stamped.header.stamp = self.get_clock().now().to_msg()
        self.map_to_odom.header.stamp = self.get_clock().now().to_msg()
        # 发布变换
        self.tf_broadcaster.sendTransform(self.map_to_odom)
        self.tf_broadcaster.sendTransform(self.transform_stamped)


def main(args=None):
    rclpy.init(args=args)
    static_tf_publisher = StaticTfPublisher()
    rclpy.spin(static_tf_publisher)
    rclpy.shutdown()

if __name__ == '__main__':
    main()